import math

import numpy as np

from Utils.transforms import transform_preds
import config as cfg


def get_max_preds(batch_heatmaps):
    '''
    从热力图获得预测结果
    get predictions from score maps
    heatmaps: numpy.ndarray([batch_size, num_joints, height, width])
    return:最大值坐标（索引），最大值
    '''
    #--------错误处理-----------
    assert isinstance(batch_heatmaps, np.ndarray),'batch_heatmaps should be numpy.ndarray'
    assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim'
    #依次获取3个维度的信息，batch size, keypoints numbers, heatmaps flattend
    batch_size = batch_heatmaps.shape[0]
    num_keypoints = batch_heatmaps.shape[1]
    width = batch_heatmaps.shape[3]
    heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_keypoints, -1)) #将后两维度热力图宽高压平
    idx = np.argmax(heatmaps_reshaped, 2)   #在axis=2中，获得最大值的索引
    maxvals = np.amax(heatmaps_reshaped, 2) #在axis=2中，获得最大值

    #扩张维度shape=2 -> shape=3
    maxvals = maxvals.reshape((batch_size, num_keypoints, 1))   #最大值
    idx = idx.reshape((batch_size, num_keypoints, 1))           #索引
    #对最大值的索引进行处理
    preds = np.tile(idx, (1, 1, 2)).astype(np.float32)  #平铺

    preds[:, :, 0] = (preds[:, :, 0]) % width
    preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)     #np.floor对输入数组进行逐元素向下取整
    #对最大值进行处理
    pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))
    pred_mask = pred_mask.astype(np.float32)

    preds *= pred_mask
    return preds, maxvals

def get_final_preds(batch_heatmaps, center, scale):
    '''
    batch_heatmaps:模型的预测结果,ndarray:(batch_size,num_points,64,64)
    center:中心点坐标,ndarray:(batch_size,2)
    scale:扩张比例,ndarray:(batch_size,2)
    return:preds,maxvals
        preds是一个ndarray:(batch_size,num_points,2)，像是坐标
        maxvals是一个ndarray:(batch_size,num_points,1)，像是最大值
    '''
    coords, maxvals = get_max_preds(batch_heatmaps)

    heatmap_height = batch_heatmaps.shape[2]
    heatmap_width = batch_heatmaps.shape[3]

    #后处理
    if cfg.post_process:
        for n in range(coords.shape[0]):            #迭代batch，也就是迭代每张图片
            for p in range(coords.shape[1]):        #迭代每个关键点
                hm = batch_heatmaps[n][p]           #第n个图片的第p个点
                px = int(math.floor(coords[n][p][0] + 0.5))
                py = int(math.floor(coords[n][p][1] + 0.5))
                if 1 < px < heatmap_width-1 and 1 < py < heatmap_height-1:
                    diff = np.array([hm[py][px+1] - hm[py][px-1],
                                     hm[py+1][px]-hm[py-1][px]])
                    coords[n][p] += np.sign(diff) * .25

    preds = coords.copy()

    # Transform back
    for i in range(coords.shape[0]):    #一个循环是一个batch，对该batch下所有的图片进行处理
        preds[i] = transform_preds(coords[i], center[i], scale[i],[heatmap_width, heatmap_height])
    return preds, maxvals
